package util

import (
	"LunovaDB/errors"
	"net"
	"sync"
)

type ConnPool struct {
	pool    *sync.Pool
	conns   []net.Conn
	mu      sync.Mutex
	maxSize int
}

func NewConnPool(maxSize int, fn func() (net.Conn, error)) *ConnPool {
	return &ConnPool{
		pool: &sync.Pool{
			New: func() interface{} {
				conn, err := fn()
				if err != nil {
					return nil
				}
				return conn
			},
		},
		maxSize: maxSize,
		conns:   make([]net.Conn, 0, maxSize),
	}
}

func (c *ConnPool) Get() (net.Conn, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if len(c.conns) > 0 {
		Listener := c.conns[len(c.conns)-1]
		c.conns = c.conns[:len(c.conns)-1]
		return Listener, nil
	}
	Conn := c.pool.Get()
	if Conn == nil {
		return nil, errors.CreatConnError
	}
	return Conn.(net.Conn), nil
}

func (c *ConnPool) Put(conns net.Conn) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if len(c.conns) < c.maxSize {
		c.conns = append(c.conns, conns)
	} else {
		conns.Close()
	}
}

func (c *ConnPool) Close() {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, Listener := range c.conns {
		Listener.Close()
	}
	c.conns = nil
}
